package org.zjx.client;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import lombok.Data;
import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.core.SocketConfig;
import org.zjx.exception.RpcTimeoutException;
import org.zjx.exception.ServiceUnavailableException;
import org.zjx.message.IMessage;
import org.zjx.message.MessageFactory;
import org.zjx.message.messages.*;

import java.io.*;
import java.lang.reflect.*;
import java.lang.reflect.Proxy;
import java.net.*;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

@Data
public class SocketClient {
    private static final Logger logger = LoggerFactory.getLogger(SocketClient.class);
    private final SocketConfig config;
    private Socket socket;
    private DataOutputStream dataOut;
    private DataInputStream dataIn;
    private ScheduledExecutorService heartbeatExecutor;
    private volatile boolean connected = false;
    private MessageListener messageListener;
    private Thread receiveThread;
    private ScheduledExecutorService reconnectScheduler;

    // 修复类型声明：使用 IMessage 而不是 RpcResponse
    private final Map<String, CompletableFuture<IMessage>> pendingRequests = new ConcurrentHashMap<>();

    private final AtomicBoolean isReconnecting = new AtomicBoolean(false);
    private final SimpleDateFormat dateFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");

    public SocketClient(SocketConfig config) {
        // 关键修复：创建配置的深拷贝，防止外部修改
        this.config = new SocketConfig()
                .setServerHost(config.getServerHost())
                .setPort(config.getPort())
                .setClientPort(config.getClientPort())
                .setHeartbeatInterval(config.getHeartbeatInterval())
                .setRequestTimeout(config.getRequestTimeout())
                .setReconnectInterval(config.getReconnectInterval())
                .setMaxReconnectAttempts(config.getMaxReconnectAttempts());
    }

    @SneakyThrows
    public void connect() {
        if (isReconnecting.get()) {
            logger.info("Connecting to server {}:{}", config.getServerHost(), config.getPort());
            return;
        }
        if (!connected) {
            try {
                initConnection();
                startHeartbeat();
                startReceiveThread();
            } catch (IOException e) {
                handleConnectionFailure(e);
            }
        }
        logger.info("Connecting to {}:{} (local port {})",
                config.getServerHost(),
                config.getPort(),
                socket.getLocalPort()
        );
    }

    /**
     * 启动消息接收线程
     */
    private void startReceiveThread() {
        if (receiveThread != null && receiveThread.isAlive()) {
            receiveThread.interrupt();
        }
        receiveThread = new Thread(this::receiveMessages, "SocketClient-Receiver");
        receiveThread.setDaemon(true);
        receiveThread.start();
    }

    /**
     * 初始化底层连接（Socket、输入输出流）
     */
    private void initConnection() throws IOException {
        // 添加客户端端口绑定逻辑
        if (config.getClientPort() > 0) {
            // 绑定指定客户端端口
            InetAddress bindAddress = InetAddress.getByName(config.getServerHost());
            SocketAddress localAddr = new InetSocketAddress(bindAddress, config.getClientPort());

            this.socket = new Socket();
            // 设置重用地址选项
            this.socket.setReuseAddress(true);
            // 绑定客户端端口
            this.socket.bind(localAddr);
            // 连接到服务器
            this.socket.connect(new InetSocketAddress(config.getServerHost(), config.getPort()));
        } else {
            // 使用默认方式创建Socket
            this.socket = new Socket(config.getServerHost(), config.getPort());
        }

        this.dataOut = new DataOutputStream(socket.getOutputStream());
        this.dataIn = new DataInputStream(socket.getInputStream());
        this.connected = true;

        logger.info("Connected to server {}:{} using local port {}",
                config.getServerHost(), config.getPort(),
                socket.getLocalPort());
    }

    /**
     * 处理连接失败（触发重连）
     */
    private void handleConnectionFailure(IOException e) {
        connected = false;
        logger.error("Connection failed: {}", e.getMessage());
        if (!isReconnecting.compareAndSet(false, true)) {
            return;
        }
        scheduleReconnect();
    }

    /**
     * 调度重连任务（指数退避策略）
     */
    private void scheduleReconnect() {
        if (reconnectScheduler != null) reconnectScheduler.shutdown();
        reconnectScheduler = Executors.newSingleThreadScheduledExecutor();

        AtomicLong initialDelay = new AtomicLong(config.getReconnectInterval());
        long maxDelay = initialDelay.get() * 2; // 最大延迟时间
        reconnectScheduler.scheduleAtFixedRate(() -> {
            isReconnecting.set(false);
            connect();
            if (connected) {
                reconnectScheduler.shutdown();
            } else {
                // 指数退避：每次延迟翻倍，直到达到最大延迟
                long nextDelay = Math.min(initialDelay.get() * 2, maxDelay);
                initialDelay.set(nextDelay);
            }
        }, initialDelay.get(), initialDelay.get(), TimeUnit.MILLISECONDS);
    }


    public IMessage sendAndWait(IMessage message, long timeoutMillis) {
        int attempts = 0;
        while (attempts < 3) {
            try {
                return doSendAndWait(message, timeoutMillis);
            } catch (Exception e) {
                attempts++;
                logger.warn("Connection reset, reconnecting (attempt {})", attempts);
                reconnect();
            }
        }
        throw new ServiceUnavailableException("Failed after 3 attempts");
    }


    private void reconnect() {
        disconnect();
        connect();
    }

    /**
     * 发送消息并等待响应
     */
    public IMessage doSendAndWait(IMessage message, long timeoutMillis) {
        // 检查连接状态
        if (!this.connected || this.dataOut == null) {
            throw new ServiceUnavailableException("Connection is not active");
        }

        CompletableFuture<IMessage> future = new CompletableFuture<>();
        String messageId = message.getMessageId();
        if (messageId == null) {
            messageId = UUID.randomUUID().toString();
            message.setMessageId(messageId);
        }

        pendingRequests.put(messageId, future);

        try {
            sendMessage(message);
            logger.debug("Sent message with ID: {}", messageId);
        } catch (ServiceUnavailableException e) {
            pendingRequests.remove(messageId);
            throw e;
        }

        try {
            return future.get(timeoutMillis, TimeUnit.MILLISECONDS);
        } catch (TimeoutException e) {
            pendingRequests.remove(messageId);
            throw new RpcTimeoutException("Request timed out after " + timeoutMillis + "ms");
        } catch (Exception e) {
            pendingRequests.remove(messageId);
            throw new RpcTimeoutException("Error while waiting for response", e);
        }
    }

    @SuppressWarnings("unchecked")
    public <T> T createProxy(Class<T> interfaceClass) {
        return (T) Proxy.newProxyInstance(
                interfaceClass.getClassLoader(),
                new Class<?>[]{interfaceClass},
                new RemoteInvocationHandler(interfaceClass, this)
        );
    }

    private class RemoteInvocationHandler implements InvocationHandler {
        private final Class<?> interfaceClass;
        private final SocketClient client;

        public RemoteInvocationHandler(Class<?> interfaceClass, SocketClient client) {
            this.interfaceClass = interfaceClass;
            this.client = client;
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            // 构建请求
            RpcRequest request = new RpcRequest();
            request.setMessageId(UUID.randomUUID().toString());
            request.setClassName(interfaceClass.getName());
            request.setMethodName(method.getName());

            // 获取方法参数类型
            Class<?>[] paramTypes = method.getParameterTypes();

            // 转换为类名字符串数组
            String[] typeNames = new String[paramTypes.length];
            for (int i = 0; i < paramTypes.length; i++) {
                typeNames[i] = paramTypes[i].getName();
            }

            // 处理参数，确保枚举类型正确传递
            Object[] processedArgs = new Object[args != null ? args.length : 0];
            for (int i = 0; i < processedArgs.length; i++) {
                Object arg = args[i];
                if (arg != null && paramTypes[i].isEnum()) {
                    processedArgs[i] = ((Enum<?>)arg).name();
                } else {
                    processedArgs[i] = arg;
                }
            }

            request.setParameterTypeNames(typeNames);
            request.setParameters(processedArgs);

            CompletableFuture<IMessage> future = new CompletableFuture<>();
            client.getPendingRequests().put(request.getMessageId(), future);

            try {
                client.sendMessage(request);
                IMessage response = future.get(config.getRequestTimeout(), TimeUnit.MILLISECONDS);

                if (response instanceof RpcResponse rpcResponse) {
                    if (rpcResponse.getError() != null) {
                        throw rpcResponse.getError();
                    }

                    Object result = rpcResponse.getResult();
                    if (result == null) return null;

                    // 通用集合类型处理
                    if (List.class.isAssignableFrom(method.getReturnType())) {
                        Type returnType = method.getGenericReturnType();
                        if (returnType instanceof ParameterizedType) {
                            Type[] actualTypeArgs = ((ParameterizedType) returnType).getActualTypeArguments();
                            if (actualTypeArgs.length > 0 && actualTypeArgs[0] instanceof Class<?> elementType) {
                                String json = result instanceof String ?
                                        (String)result : JSON.toJSONString(result);
                                return JSON.parseArray(json, elementType);
                            }
                        }
                    }

                    // 强制类型转换保障
                    if (result instanceof JSONObject jsonObj) {
                        result = JSON.toJavaObject(jsonObj, method.getReturnType());
                    }

                    return method.getReturnType().cast(result);
                }
                throw new RuntimeException("Unexpected response type: " + response.getClass().getName());
            } catch (TimeoutException e) {
                future.completeExceptionally(new RpcTimeoutException("Request timed out"));
                throw new RpcTimeoutException("Request timed out after " + config.getRequestTimeout() + "ms");
            } catch (ServiceUnavailableException e) {
                future.completeExceptionally(e);
                client.getPendingRequests().remove(request.getMessageId(), future);
                throw new ServiceUnavailableException("Service unavailable: connection failed", e);
            }catch (ExecutionException e) {
                Throwable cause = e.getCause();
                throw cause != null ? cause : e;
            }
        }
    }

    public void disconnect() {
        if (reconnectScheduler != null) reconnectScheduler.shutdownNow();
        if (this.connected) {
            this.connected = false;
            try {
                // 1. 关闭输出流（触发服务端读取结束）
                if (this.dataOut != null) {
                    this.dataOut.close();
                    this.dataOut = null;
                }

                // 2. 停止心跳线程（等待当前任务完成）
                if (this.heartbeatExecutor != null) {
                    this.heartbeatExecutor.shutdown();
                    if (!this.heartbeatExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                        this.heartbeatExecutor.shutdownNow();
                    }
                }

                // 3. 中断接收线程并等待
                if (this.receiveThread != null && this.receiveThread.isAlive()) {
                    this.receiveThread.interrupt();
                    this.receiveThread.join(5000); // 等待最多5秒
                }

                // 4. 关闭输入流（此时应已无阻塞）
                if (this.dataIn != null) {
                    this.dataIn.close();
                    this.dataIn = null;
                }

                // 5. 最后关闭套接字
                if (this.socket != null) {
                    this.socket.close();
                    this.socket = null;
                }

                logger.info("Disconnecting from server");
            } catch (Exception e) {
                logger.error("Error closing connection", e);
            }
        }
    }

    public void sendCommand(String command, Map<String, Object> params) {
        if (!this.connected) {
            logger.warn("Not connected, command not sent");
        } else {
            CommandMessage cmd = new CommandMessage(command);
            cmd.setMessageId("CMD_" + System.currentTimeMillis());
            cmd.setParams(params);
            this.sendMessage(cmd);
        }
    }

    public void sendMessage(IMessage message) throws ServiceUnavailableException {
        if (!this.connected || this.dataOut == null) {
            throw new ServiceUnavailableException("Connection is not active");
        }

        try {
            byte[] data = MessageFactory.serialize(message);

            // 添加调试日志
            if (logger.isDebugEnabled()) {
                logger.debug("Sending message ({} bytes): {}", data.length,
                        new String(data, StandardCharsets.UTF_8));
            }

            // 先写入长度头
            dataOut.writeInt(data.length);
            // 再写入数据
            dataOut.write(data);
            dataOut.flush();
        } catch (Exception e) {
            this.connected = false;
            logger.error("Failed to send message ({} bytes)",
                    message != null ? MessageFactory.serialize(message).length : "null", e);
            throw new ServiceUnavailableException("Failed to send message", e);
        }
    }

    /**
     * 启动心跳线程
     */
    private void startHeartbeat() {
        // 添加非正数心跳间隔检查
        if (config.getHeartbeatInterval() <= 0) {
            logger.debug("Heartbeat disabled for connection to {}:{}",
                    config.getServerHost(), config.getPort());
            return; // 直接返回，不启动心跳
        }

        if (heartbeatExecutor != null) {
            heartbeatExecutor.shutdownNow();
        }
        heartbeatExecutor = Executors.newSingleThreadScheduledExecutor();
        heartbeatExecutor.scheduleAtFixedRate(() -> {
            if (connected) {
                try {
                    sendHeartbeat();
                } catch (Exception e) {
                    logger.error("Heartbeat failed", e);
                    disconnect();
                }
            }
        }, 0, config.getHeartbeatInterval(), TimeUnit.MILLISECONDS);
    }

    /**
     * 发送心跳消息
     */
    private void sendHeartbeat() {
        HeartbeatMessage heartbeat = new HeartbeatMessage();
        heartbeat.setMessageId("HB_" + System.currentTimeMillis());
        sendMessage(heartbeat);
        logger.debug("Heartbeat sent");
    }

    private void receiveMessages() {
        try {
            while (!Thread.currentThread().isInterrupted() && this.connected) {
                try {
                    // 使用新的读取方法
                    int length = readLengthHeader(dataIn);
                    byte[] data = readMessageBody(dataIn, length);

                    // 反序列化消息
                    IMessage message = MessageFactory.safeDeserialize(data);

                    // 处理消息
                    processMessage(message);

                } catch (SocketTimeoutException e) {
                    handleReceiveTimeout();
                } catch (EOFException e) {
                    logger.info("Server closed the connection gracefully");
                    disconnect();
                    break;
                } catch (SocketException e) {
                    if (this.connected) {
                        logger.info("Socket error: {}", e.getMessage());
                        disconnect();
                    }
                    break;
                }
            }
        } catch (Exception e) {
            if (this.connected) {
                logger.error("Unexpected error in receive thread", e);
                disconnect();
            }
        } finally {
            logger.debug("Receive thread exiting");
        }
    }

    /**
     * 读取消息长度头（带超时和校验）
     */
    private int readLengthHeader(DataInputStream in) throws IOException {
        long startTime = System.currentTimeMillis();
        final int headerSize = 4; // int类型的长度头占4字节

        // 等待直到有足够数据或超时
        while (in.available() < headerSize) {
            // 检查连接状态
            if (!this.connected || socket.isClosed()) {
                throw new SocketException("Connection is not active");
            }

            // 检查超时
            if (System.currentTimeMillis() - startTime > config.getRequestTimeout()) {
                throw new SocketTimeoutException("Timeout waiting for length header");
            }

            // 避免忙等待
            try {
                Thread.sleep(10);
            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
                throw new IOException("Interrupted while reading length header");
            }
        }

        // 读取长度头
        int length = in.readInt();

        // 验证长度合理性
        if (length <= 0) {
            throw new IOException("Invalid message length: " + length + " (<= 0)");
        }

        // 设置合理的最大消息大小限制（例如10MB）
        if (length > 10 * 1024 * 1024) {
            throw new IOException("Message too large: " + length + " bytes");
        }

        return length;
    }

    /**
     * 读取消息体内容
     */
    private byte[] readMessageBody(DataInputStream in, int length) throws IOException {
        byte[] data = new byte[length];
        int totalRead = 0;
        long startTime = System.currentTimeMillis();

        while (totalRead < length) {
            // 检查连接状态
            if (!this.connected || socket.isClosed()) {
                throw new SocketException("Connection closed while reading");
            }

            // 检查超时
            if (System.currentTimeMillis() - startTime > config.getRequestTimeout()) {
                throw new SocketTimeoutException("Timeout reading message body");
            }

            // 读取数据
            int read = in.read(data, totalRead, length - totalRead);
            if (read == -1) {
                throw new EOFException("Unexpected end of stream");
            }
            totalRead += read;

            // 添加调试日志
            if (logger.isTraceEnabled() && totalRead < length) {
                logger.trace("Partial read: {}/{} bytes received", totalRead, length);
            }
        }

        if (logger.isDebugEnabled()) {
            logger.debug("Completed reading message body: {} bytes", length);
        }

        return data;
    }

    /**
     * 处理接收超时
     */
    private void handleReceiveTimeout() {
        if (this.connected) {
            logger.warn("Receive timeout occurred, checking connection...");

            // 检查物理连接状态
            boolean isActuallyConnected = socket != null
                    && !socket.isClosed()
                    && socket.isConnected();

            if (!isActuallyConnected) {
                logger.info("Connection is no longer active, disconnecting...");
                disconnect();
            } else {
                logger.debug("Connection is still active, may be temporary network issue");
            }
        }
    }

    private void handleReceiveError(IOException e) {
        if (this.connected) {
            logger.error("Receive error", e);
            disconnect();
        }
    }

    /**
     * 处理单条消息（根据类型分发）
     */
    private void processMessage(IMessage message) {
        try {
            String messageId = message.getMessageId();

            // 处理无 messageId 的消息
            if (messageId == null) {
                if (messageListener != null) {
                    messageListener.onMessage(message);
                } else if (message instanceof HeartbeatMessage) {
                    logger.debug("Received heartbeat without ID");
                }
                return;
            }

            // 处理等待中的请求
            CompletableFuture<IMessage> future = pendingRequests.get(messageId);
            if (future != null) {
                future.complete(message);
                pendingRequests.remove(messageId);
                return;
            }

            // 处理特定消息类型
            if (message instanceof HeartbeatMessage) {
                logger.debug("Received heartbeat");
            } else if (messageListener != null) {
                messageListener.onMessage(message);
            }
        } catch (Exception e) {
            logger.error("Failed to process message: " + message, e);
        }
    }

    public interface MessageListener {
        void onMessage(IMessage message);
    }
}